from gymnasium import spaces
from gymnasium.envs.mujoco import MujocoEnv

import mujoco

import numpy as np
from pathlib import Path

# 默认相机配置
DEFAULT_CAMERA_CONFIG = {
    "azimuth": 90.0,
    "distance": 2.0,
    "elevation": -25.0,
    "lookat": np.array([0.0, 0.0, 0.0]),
    "fixedcamid": 0,
    "trackbodyid": -1,
    "type": 2,
}


class BittleMujocoEnv(MujocoEnv):
    """遵循gym接口的自定义环境。"""

    metadata = {
        "render_modes": [
            "human",
            "rgb_array",
            "depth_array",
        ],
    }

    def __init__(self, ctrl_type="position", **kwargs):
        model_path = Path(f"./bittle/scene_{ctrl_type}.xml")
        MujocoEnv.__init__(
            self,
            model_path=model_path.absolute().as_posix(),
            frame_skip=10,  # 每10帧执行一次动作
            observation_space=None,  # 之后手动设置
            default_camera_config=DEFAULT_CAMERA_CONFIG,
            **kwargs,
        )

        # 更新元数据以包含渲染FPS
        self.metadata = {
            "render_modes": [
                "human",
                "rgb_array",
                "depth_array",
            ],
            "render_fps": 60,
        }
        self._last_render_time = -1.0
        self._max_episode_time_sec = 15.0  # 8秒
        self._step = 0

        # 奖励和成本函数的权重
        self.reward_weights = {
            "linear_vel_tracking": 30,
            "angular_vel_tracking": 80,
            "healthy": 40.0,
            "feet_airtime": -5,
        }
        self.cost_weights = {
            "torque": 0.0002,
            "vertical_vel": 1.0,  # z轴线速度惩罚
            "xy_angular_vel": 80,  # 偏航角的角速度惩罚
            "action_rate": -30,
            "joint_limit": 70,
            # "knee_height": 40, # 膝关节高度成本,新增
            "joint_velocity": -10,
            "joint_acceleration": 0,
            "orientation": 1.0,
            "collision": 0.8,
            "default_joint_position": 30,
            "unhealthy": 50.0,
        }

        self._curriculum_base = 0.3
        self._gravity_vector = np.array(self.model.opt.gravity)
        self._default_joint_position = np.array(self.model.key_ctrl[0])

        # 坐标轴示意图
        #        ^ y (前方)
        #      __|___
        #     |      |
        #     |      |
        #     |      |--> x (右侧)
        #     |______|
        #

        # vx (m/s), vy (m/s), wz (rad/s), y是机器人向前
        self._desired_velocity_min = np.array([0.0, -0.5, 0.0])
        self._desired_velocity_max = np.array([0.0, -0.5, 0.0])
        self._desired_velocity = self._sample_desired_vel()  # [0.0, 0.3, 0.0]
        self._obs_scale = {
            "linear_velocity": 2.0,
            "angular_velocity": 0.25,
            "dofs_position": 1.0,
            "dofs_velocity": 0.05,
        }
        self._tracking_velocity_sigma = 0.05  # 速度跟踪奖励的标准差

        # 用于确定是否应终止回合的指标
        self._healthy_z_range = (0.05, 0.09)  # 6-9cm, 初始7.5cm
        self._healthy_knee_height_range = (0.02, 0.08)  # 膝关节允许的高度范围(2-8cm)
        self._healthy_pitch_range = (-np.deg2rad(10), np.deg2rad(10))
        self._healthy_roll_range = (-np.deg2rad(10), np.deg2rad(10))

        self._feet_air_time = np.zeros(4)
        self._last_contacts = np.zeros(4)
        self._cfrc_ext_feet_indices = [2, 4, 6, 8]  # 2:FL, 4:FR, 6:RL, 8:RR
        self._cfrc_ext_contact_indices = [3, 5, 7, 9]

        # 控制关节的非惩罚自由度范围
        dof_position_limit_multiplier = 0.5  # 不受惩罚的范围百分比
        ctrl_range_offset = (
            0.5
            * (1 - dof_position_limit_multiplier)
            * (
                self.model.actuator_ctrlrange[:, 1]  # 解释
                - self.model.actuator_ctrlrange[:, 0]
            )
        )
        # 第一个值是根关节，所以我们忽略它
        self._soft_joint_range = np.copy(self.model.actuator_ctrlrange)
        self._soft_joint_range[:, 0] += ctrl_range_offset
        self._soft_joint_range[:, 1] -= ctrl_range_offset

        self._reset_noise_scale = 0.08

        # 动作：8个关节的角度
        self._last_action = np.zeros(8)

        self._clip_obs_threshold = 100.0
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=self._get_obs().shape, dtype=np.float32
        )

        self.action_space = spaces.Box(
            low=self.model.actuator_ctrlrange[:, 0],
            high=self.model.actuator_ctrlrange[:, 1],
            shape=(8,),
            dtype=np.float32,
        )

        # 足部站点名称到索引的映射
        # https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-site
        # https://mujoco.readthedocs.io/en/stable/APIreference/APItypes.html#mjtobj
        feet_site = [
            "rfk",
            "lfk",
            "rbk",
            "lbk",
        ]
        self._feet_site_name_to_id = {
            f: mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE.value, f)
            for f in feet_site
        }

        # 膝盖站点名称到索引的映射
        knee_site = [
            "rfk",
            "lfk",
            "rbk",
            "lbk",
        ]
        self._knee_site_name_to_id = {
            f: mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE.value, f)
            for f in knee_site
        }

        self._main_body_id = mujoco.mj_name2id(
            self.model, mujoco.mjtObj.mjOBJ_BODY.value, "torso"
        )

    def step(self, action):
        self._step += 1
        self.do_simulation(action, self.frame_skip)

        observation = self._get_obs()
        reward, reward_info = self._calc_reward(action)
        # TODO: 考虑在膝盖触地时终止
        terminated = not self.is_healthy
        truncated = self._step >= (self._max_episode_time_sec / self.dt)
        info = {
            "x_position": self.data.qpos[0],
            "y_position": self.data.qpos[1],
            "distance_from_origin": np.linalg.norm(self.data.qpos[0:2], ord=2),
            **reward_info,
        }

        if self.render_mode == "human" and (self.data.time - self._last_render_time) > (
            1.0 / self.metadata["render_fps"]
        ):
            self.render()
            self._last_render_time = self.data.time

        self._last_action = action

        return observation, reward, terminated, truncated, info

    @property
    def is_healthy(self):
        state = self.state_vector()
        min_z, max_z = self._healthy_z_range
        is_healthy = np.isfinite(state).all() and min_z <= state[2] <= max_z

        min_roll, max_roll = self._healthy_roll_range
        is_healthy = is_healthy and min_roll <= state[4] <= max_roll

        min_pitch, max_pitch = self._healthy_pitch_range
        is_healthy = is_healthy and min_pitch <= state[5] <= max_pitch

        # min_knee_height, max_knee_height = self._healthy_knee_height_range
        # for(knee_site_id) in self._knee_site_name_to_id.items():
        #     knee_height = self.data.site_xpos[knee_site_id][2]
        #     is_healthy = is_healthy and min_knee_height <= knee_height <= max_knee_height

        if state[2] < min_z:
            print(f"Terminated: Z position {state[2]} < {min_z}")
        elif state[2] > max_z:
            print(f"Terminated: Z position {state[2]} > {max_z}")
        elif state[4] < min_roll:
            print(f"Terminated: Roll {state[4]} < {min_roll}")
        elif state[4] > max_roll:
            print(f"Terminated: Roll {state[4]} > {max_roll}")
        elif state[5] < min_pitch:
            print(f"Terminated: Pitch {state[5]} < {min_pitch}")
        elif state[5] > max_pitch:
            print(f"Terminated: Pitch {state[5]} > {max_pitch}")
            # 如果膝关节高度不合适，打印提示信息
        # if knee_height < min_knee_height:
        #     print(f"Terminated: {knee_site_name} knee too low: {knee_height:.3f} < {min_knee_height}")
        # elif knee_height > max_knee_height:
        #     print(f"Terminated: {knee_site_name} knee too high: {knee_height:.3f} > {max_knee_height}")

        return is_healthy
        # return True

    @property
    def projected_gravity(self):
        w, x, y, z = self.data.qpos[3:7]
        euler_orientation = np.array(self.euler_from_quaternion(w, x, y, z))
        projected_gravity_not_normalized = (
            np.dot(self._gravity_vector, euler_orientation) * euler_orientation
        )
        if np.linalg.norm(projected_gravity_not_normalized) == 0:
            return projected_gravity_not_normalized
        else:
            return projected_gravity_not_normalized / np.linalg.norm(
                projected_gravity_not_normalized
            )

    @property
    def feet_contact_forces(self):
        feet_contact_forces = self.data.cfrc_ext[self._cfrc_ext_feet_indices]
        return np.linalg.norm(feet_contact_forces, axis=1)

    ######### 正奖励函数 #########
    @property
    def linear_velocity_tracking_reward(self):
        vel_sqr_error = np.sum(
            np.square(self._desired_velocity[:2] - self.data.qvel[:2])
        )
        # print(f"realtime velocity: {self.data.qvel[:2]}")
        # print(f"desired velocity: {self._desired_velocity[:2]}")
        # print(f"realtime location: {self.data.qpos[:2]}")
        return np.exp(-vel_sqr_error / self._tracking_velocity_sigma)

    @property
    def angular_velocity_tracking_reward(self):
        vel_sqr_error = np.square(self._desired_velocity[2] - self.data.qvel[5])
        return np.exp(-vel_sqr_error / self._tracking_velocity_sigma)

    @property
    def heading_tracking_reward(self):
        # TODO: qpos[3:7] 是四元数值
        pass

    @property
    def feet_air_time_reward(self):
        """根据步幅的持续时间奖励步幅，仅当足部接触地面时"""
        feet_contact_force_mag = self.feet_contact_forces
        curr_contact = feet_contact_force_mag > 1.0
        contact_filter = np.logical_or(curr_contact, self._last_contacts)
        self._last_contacts = curr_contact

        # 如果 feet_air_time > 0（足部在空中）并且 contact_filter 检测到与地面的接触
        # 那么这是该步幅的第一次接触
        first_contact = (self._feet_air_time > 0.0) * contact_filter
        self._feet_air_time += self.dt

        # 奖励刚刚完成步幅的足部（第一次接触）
        air_time_reward = np.sum((self._feet_air_time - 1.0) * first_contact)
        # 如果期望速度非常低，则不奖励（即机器人应保持静止，足部不应移动）
        air_time_reward *= np.linalg.norm(self._desired_velocity[:2]) > 0.1

        # 将刚刚接触地面的足部的空中时间清零（即 contact_filter==1）
        self._feet_air_time *= ~contact_filter

        return air_time_reward

    @property
    def healthy_reward(self):
        return self.is_healthy

    @property
    # def longer_stride_reward(self, action):
    #     # 鼓励更长的步幅
    #     return (self._last_action - action)

    ######### 负奖励函数 #########
    @property  # TODO: 未使用
    def feet_contact_forces_cost(self):
        return np.sum(
            (self.feet_contact_forces - self._max_contact_force).clip(min=0.0)
        )

    @property
    def unhelthy_cost(self):
        # 惩罚机器姿态不健康
        return not self.is_healthy

    @property
    def non_flat_base_cost(self):
        # 惩罚机器人不平放在地面上
        return np.sum(np.square(self.projected_gravity[:2]))

    @property
    def collision_cost(self):
        # 惩罚选定身体的碰撞
        return np.sum(
            1.0
            * (np.linalg.norm(self.data.cfrc_ext[self._cfrc_ext_contact_indices]) > 0.1)
        )

    @property
    def joint_limit_cost(self):
        # 惩罚机器人关节超出软控制范围
        out_of_range = (self._soft_joint_range[:, 0] - self.data.qpos[7:]).clip(
            min=0.0
        ) + (self.data.qpos[7:] - self._soft_joint_range[:, 1]).clip(min=0.0)
        return np.sum(out_of_range)

    @property
    def knee_height_cost(self):
        """计算膝关节高度成本"""
        min_knee_height, max_knee_height = self._healthy_knee_height_range
        ideal_height = (min_knee_height + max_knee_height) / 2  # 理想高度为范围中点
        total_cost = 0.0

        for knee_site_id in self._knee_site_name_to_id.values():
            knee_height = self.data.site_xpos[knee_site_id][2]
            # 相对于理想高度的平方误差
            cost = np.square(knee_height - ideal_height)
            total_cost += cost

        return total_cost

    @property
    def torque_cost(self):
        # 最后8个值是电机扭矩
        return np.sum(np.square(self.data.qfrc_actuator[-8:]))

    @property
    def vertical_velocity_cost(self):
        return np.square(self.data.qvel[2])

    @property
    def xy_angular_velocity_cost(self):
        return np.sum(np.square(self.data.qvel[3:5]))

    def action_rate_cost(self, action):
        return np.sum(np.square(self._last_action - action))

    @property
    def joint_velocity_cost(self):
        return np.sum(np.square(self.data.qvel[6:]))

    @property
    def acceleration_cost(self):
        return np.sum(np.square(self.data.qacc[6:]))

    @property
    def default_joint_position_cost(self):
        return np.sum(np.square(self.data.qpos[7:] - self._default_joint_position))

    @property
    def smoothness_cost(self):
        return np.sum(np.square(self.data.qpos[7:] - self._last_action))

    @property
    def curriculum_factor(self):
        return self._curriculum_base**0.997

    def _calc_reward(self, action):
        # TODO: 添加调试模式，使用自定义Tensorboard调用单个奖励函数，以更好地了解每个奖励函数的贡献
        # TODO: 大腿或小腿接触地面的成本

        # 正奖励
        linear_vel_tracking_reward = (
            self.linear_velocity_tracking_reward
            * self.reward_weights["linear_vel_tracking"]
        )
        angular_vel_tracking_reward = (
            self.angular_velocity_tracking_reward
            * self.reward_weights["angular_vel_tracking"]
        )
        healthy_reward = self.healthy_reward * self.reward_weights["healthy"]
        feet_air_time_reward = (
            self.feet_air_time_reward * self.reward_weights["feet_airtime"]
        )
        rewards = (
            linear_vel_tracking_reward
            + angular_vel_tracking_reward
            + healthy_reward
            + feet_air_time_reward
        )

        # 负成本
        ctrl_cost = self.torque_cost * self.cost_weights["torque"]
        action_rate_cost = (
            self.action_rate_cost(action) * self.cost_weights["action_rate"]
        )
        vertical_vel_cost = (
            self.vertical_velocity_cost * self.cost_weights["vertical_vel"]
        )
        xy_angular_vel_cost = (
            self.xy_angular_velocity_cost * self.cost_weights["xy_angular_vel"]
        )
        joint_limit_cost = self.joint_limit_cost * self.cost_weights["joint_limit"]
        # knee_height_cost = self.knee_height_cost * self.cost_weights["knee_height"]
        joint_velocity_cost = (
            self.joint_velocity_cost * self.cost_weights["joint_velocity"]
        )
        joint_acceleration_cost = (
            self.acceleration_cost * self.cost_weights["joint_acceleration"]
        )
        orientation_cost = self.non_flat_base_cost * self.cost_weights["orientation"]
        collision_cost = (
            self.collision_cost * self.cost_weights["collision"]
        )  # TODO: 未使用
        default_joint_position_cost = (
            self.default_joint_position_cost
            * self.cost_weights["default_joint_position"]
        )
        body_unhealthy_cost = self.unhelthy_cost * self.cost_weights["unhealthy"]
        costs = (
            ctrl_cost
            + action_rate_cost
            + vertical_vel_cost
            + xy_angular_vel_cost
            + joint_limit_cost
            # + knee_height_cost
            + joint_velocity_cost
            + joint_acceleration_cost
            + orientation_cost
            + default_joint_position_cost
            + body_unhealthy_cost
        )

        # reward = rewards - costs
        reward = rewards - self.curriculum_factor * costs
        reward_info = {
            "linear_vel_tracking_reward": linear_vel_tracking_reward,
            "reward_ctrl": -ctrl_cost,
            "reward_survive": healthy_reward,
        }

        return reward, reward_info

    def _get_obs(self):
        # 位置信息
        # [0,1,2]是机器人躯干的全局x,y,z位置
        # [3,4,5,6]是表示机器人方向的四元数
        # 上述七个值被忽略，因为它们是上帝信息
        # 剩下的8个值是关节位置
        # 关节位置相对于起始位置，8维
        dofs_position = self.data.qpos[7:].flatten() - self.model.key_qpos[0, 7:]

        velocity = self.data.qvel.flatten()
        # [0,1,2]是机器人的全局线速度
        # [3,4,5]是机器人的角速度
        # 剩下的8个值是关节速度
        base_linear_velocity = velocity[:3]  # 3维
        base_angular_velocity = velocity[3:6]  # 3维
        dofs_velocity = velocity[6:]  # 8自由度的关节速度,8维

        # 机器人的期望速度,3维
        desired_vel = self._desired_velocity
        # 上一步的动作,8维
        last_action = self._last_action
        # 机器人的投影重力,3维
        projected_gravity = self.projected_gravity

        # 观测空间共：3+3+3+3+8+8+8=26维
        curr_obs = np.concatenate(
            (
                base_linear_velocity * self._obs_scale["linear_velocity"],
                base_angular_velocity * self._obs_scale["angular_velocity"],
                projected_gravity,
                desired_vel * self._obs_scale["linear_velocity"],
                dofs_position * self._obs_scale["dofs_position"],
                dofs_velocity * self._obs_scale["dofs_velocity"],
                last_action,
            )
        ).clip(-self._clip_obs_threshold, self._clip_obs_threshold)

        return curr_obs

    # def reset_model(self):
    #     # 使用噪声重置位置和控制值
    #     self.data.qpos[:] = self.model.key_qpos[0] + self.np_random.uniform(
    #         low=-self._reset_noise_scale,
    #         high=self._reset_noise_scale,
    #         size=self.model.nq,
    #     )
    #     self.data.ctrl[:] = self.model.key_ctrl[0] + self._reset_noise_scale * self.np_random.standard_normal(
    #         *self.data.ctrl.shape
    #     )

    #     # 重置变量并采样新的期望速度
    #     # self.data.qpos[:] = self.model.key_qpos[0]
    #     # self.data.ctrl[:] = self.model.key_ctrl[0]
    #     self._desired_velocity = self._sample_desired_vel()
    #     self._step = 0
    #     self._last_action = np.zeros(8)
    #     self._feet_air_time = np.zeros(4)
    #     self._last_contacts = np.zeros(4)
    #     self._last_render_time = -1.0

    #     observation = self._get_obs()

    #     return observation

    def reset_model(self):

        # 获取初始状态,qpos[0:3]是位置,qpos[3:7]是四元数,qpos[7:]是关节角度
        init_qpos = self.model.key_qpos[0].copy()

        # 关节角度添加噪声（跳过根节点的7个自由度）
        noise_joint = self.np_random.uniform(
            low=-self._reset_noise_scale,
            high=self._reset_noise_scale,
            size=self.model.nq - 7,  # 只对关节添加噪声
        )

        # 添加噪声到位置(z保持不变)
        noise_pos = self.np_random.uniform(
            low=-self._reset_noise_scale,
            high=self._reset_noise_scale,
            size=2,  # 只改变 x 和 y
        )

        # 添加噪声到四元数
        noise_quat = self.np_random.uniform(
            low=-self._reset_noise_scale, high=self._reset_noise_scale, size=4
        )

        # 关节角度添加噪声
        self.data.qpos[7:] = init_qpos[7:] + noise_joint  # 关节角度添加噪声
        self.data.qpos[:2] = init_qpos[:2] + noise_pos  # 更新 x, y 位置
        self.data.qpos[2] = init_qpos[2]
        self.data.qpos[3:7] = init_qpos[3:7] + noise_quat  # 更新四元数

        # 控制器噪声保持不变
        self.data.ctrl[:] = self.model.key_ctrl[
            0
        ] + self._reset_noise_scale * self.np_random.standard_normal(
            *self.data.ctrl.shape
        )

        # 重置其他变量
        self._desired_velocity = self._sample_desired_vel()
        self._step = 0
        self._last_action = np.zeros(8)
        self._feet_air_time = np.zeros(4)
        self._last_contacts = np.zeros(4)
        self._last_render_time = -1.0

        observation = self._get_obs()
        return observation

    def _get_reset_info(self):
        return {
            "x_position": self.data.qpos[0],
            "y_position": self.data.qpos[1],
            "distance_from_origin": np.linalg.norm(self.data.qpos[0:2], ord=2),
        }

    # 随机化初始速度
    def _sample_desired_vel(self):
        desired_vel = np.random.default_rng().uniform(
            low=self._desired_velocity_min, high=self._desired_velocity_max
        )
        return desired_vel

    @staticmethod
    def euler_from_quaternion(w, x, y, z):
        """
        将四元数转换为欧拉角（滚转、俯仰、偏航）
        滚转是绕x轴的旋转（逆时针）
        俯仰是绕y轴的旋转（逆时针）
        偏航是绕z轴的旋转（逆时针）
        """
        t0 = +2.0 * (w * x + y * z)
        t1 = +1.0 - 2.0 * (x * x + y * y)
        pitch_x = np.arctan2(t0, t1)

        t2 = +2.0 * (w * y - z * x)
        t2 = +1.0 if t2 > +1.0 else t2
        t2 = -1.0 if t2 < -1.0 else t2
        roll_y = np.arcsin(t2)

        t3 = +2.0 * (w * z + x * y)
        t4 = +1.0 - 2.0 * (y * y + z * z)
        yaw_z = np.arctan2(t3, t4)

        return roll_y, -pitch_x, yaw_z  # 以弧度为单位

    # @staticmethod
    # def euler_from_quaternion(w, x, y, z):
    #     """
    #     将四元数转换为欧拉角（滚转、俯仰、偏航）
    #     机器人坐标系,x轴为右,y轴为前,z轴为上。
    #     滚转是绕y轴的旋转(逆时针)
    #     俯仰是绕x轴的旋转(逆时针)
    #     偏航是绕z轴的旋转(逆时针)
    #     """
    #     # 滚转（Roll）绕Y轴旋转
    #     t0 = +2.0 * (w * y - z * x)
    #     t1 = +1.0 - 2.0 * (y * y + z * z)
    #     roll_y = np.arctan2(t0, t1)

    #     # 俯仰（Pitch）绕X轴旋转
    #     t2 = +2.0 * (w * x + y * z)
    #     t2 = +1.0 if t2 > +1.0 else t2
    #     t2 = -1.0 if t2 < -1.0 else t2
    #     pitch_x = np.arcsin(t2)

    #     # 偏航（Yaw）绕Z轴旋转
    #     t3 = +2.0 * (w * z - x * y)
    #     t4 = +1.0 - 2.0 * (z * z + x * x)
    #     yaw_z = np.arctan2(t3, t4)

    #     return roll_y, pitch_x, yaw_z  # 以弧度为单位
